package com.llmops.core.node;

import com.llmops.core.Context;
import com.llmops.core.Node;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.SystemMessage;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.service.V;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.weaviate.WeaviateEmbeddingStore;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class RagNode extends Node {
    interface Assistant {
        @SystemMessage("{{systemPrompt}}")
        String chat(@MemoryId int memoryId, @UserMessage String userMessage,@V("systemPrompt")String systemPrompt);
    }
    private String input;
    private String apiUrl;
    private String apiKey;
    private String model;
    private Integer historyRecordCount;

    public RagNode(String id,String title, Map<String, Object> config) {
        super(id, title,config);
        this.apiUrl = (String) config.getOrDefault("apiUrl", "https://api.deepseek.com/v1");
        this.apiKey = (String) config.getOrDefault("apiKey","sk-0c3e6a9f9f7a4e63bb855290c544183c");
        this.model = (String) config.getOrDefault("model", "deepseek-chat");
        this.historyRecordCount = (Integer) config.getOrDefault("historyRecordCount", 10);
        this.input = (String) config.get("input");
    }

    @Override
    public void execute(Context context) {
        try {
            EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
            String weaviateUrl = "76okjqrmqeat2g9ln6amg.c0.asia-southeast1.gcp.weaviate.cloud";
            String weaviateKey = "AjzWod28BlPwN5MdK7iVoIOstk1aThcl3jLR";
            EmbeddingStore<TextSegment> embeddingStore = WeaviateEmbeddingStore.builder()
                    .scheme("https").host(weaviateUrl).apiKey(weaviateKey).avoidDups(true).consistencyLevel("ALL").build();

            Embedding queryEmbedding = embeddingModel.embed(context.fillVariable(context,input)).content();
            EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                    .queryEmbedding(queryEmbedding)
                    .maxResults(2)
                    .build();
            List<EmbeddingMatch<TextSegment>> matches = embeddingStore.search(embeddingSearchRequest).matches();

            //重排序，通过大模型判断召回的文档与问题的相关性 TODO

            List<String> contents = new ArrayList<>();
            for (EmbeddingMatch<TextSegment> match : matches) {
                contents.add(match.embedded().text());
            }
            context.setVariable(id + "_output", String.join("\n",contents));
            logger.info(getId() + "_output={}", String.join("\n",contents));
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(String.format("节点={},运行异常",id));
        }
    }

    @Override
    public String getType() {
        return "ragNode";
    }
}